from fisher.agent import DQN
from fisher.models import FishNet
from fisher.environment import *
import torch
import argparse
import os
from utils.render import *

parser = argparse.ArgumentParser(description='通过DQN对原神自动钓鱼进行仿真训练')
parser.add_argument('--batch_size', default=32, type=int)
parser.add_argument('--n_states', default=3, type=int)
parser.add_argument('--n_actions', default=2, type=int)
parser.add_argument('--step_tick', default=12, type=int)
parser.add_argument('--n_episode', default=400, type=int)
parser.add_argument('--save_dir', default='./DQN_output', type=str)
parser.add_argument('--resume', default=None, type=str)
args = parser.parse_args()

if not os.path.exists(args.save_dir):
    os.makedirs(args.save_dir)
# 要被训练的强化学习网络，输入默认为3通道，输出默认为2通道，中间层铺设了16通道
net = FishNet(in_ch=args.n_states, out_ch=args.n_actions)
if args.resume:  # 读取之前保存的训练结果（因为这里是率先进行的仿真训练，所以没有读取的训练结果）
    # load_state_dict是 Module类中的方法，作用是载入其他同大小模型的训练结果
    net.load_state_dict(torch.load(args.resume))

# 定义了强化学习的参数，net是输入的要被训练的网络
agent = DQN(net, args.batch_size, args.n_states, args.n_actions, memory_capacity=2000)
# 这里定义了原神钓鱼的仿真环境，实质就是两条边缘线的随机移动
env = Fishing_sim(step_tick=args.step_tick, drawer=PltRender())
if __name__ == '__main__':
    # Start training
    print("\nCollecting experience...")
    net.train()
    for i_episode in range(args.n_episode):
        # keyboard.wait('r')
        # play 400 episodes of cartpole game
        s = env.reset()
        ep_r = 0
        while True:  # 一次死循环，跳出条件是单次的仿真时间到头或者反馈极其的低
            # 200帧后，每20帧显示一次
            if i_episode > 200 and i_episode % 20 == 0:  # 隔若干个周期展示一次学习效果
                env.render()
            # s是两个边缘位置和当前指针的位置，输入到greedy算法函数中进行当前阶段是否需要点击鼠标
            a = agent.choose_action(s)
            # 如果a=0，那么不点击鼠标，如果a=1，那么点击鼠标，通过仿真环境来模拟效果
            s_, r, done = env.step(a)  # 反馈r是得分的差值，当前得分比上次得分高，r才是正数

            # 存放每个step结束后的状态参数（位置信息、决策值、反馈还有决策后的位置信息）
            agent.store_transition(s, a, r, s_)

            ep_r += r  # 反馈积分
            # 如果经验重播缓冲区已经满了，就开始DQN学习
            #
            if agent.memory_counter > agent.memory_capacity:
                agent.train_step()
                if done:
                    print('训练次数 Ep: ', i_episode, ' |', '反馈积分 Ep_r: ', round(ep_r, 2))

            if done:
                # if game is over, then skip the while loop.
                break
            # use next state to update the current state.
            s = s_
    torch.save(net.state_dict(), os.path.join(args.save_dir, f'fish_sim_net.pth'))
